package com.sugon.xuanyuan;

import cn.smartjavaai.common.enums.DeviceEnum;
import cn.smartjavaai.ocr.config.DirectionModelConfig;
import cn.smartjavaai.ocr.config.OcrDetModelConfig;
import cn.smartjavaai.ocr.config.OcrRecModelConfig;
import cn.smartjavaai.ocr.config.OcrRecOptions;
import cn.smartjavaai.ocr.entity.OcrInfo;
import cn.smartjavaai.ocr.enums.CommonDetModelEnum;
import cn.smartjavaai.ocr.enums.CommonRecModelEnum;
import cn.smartjavaai.ocr.enums.DirectionModelEnum;
import cn.smartjavaai.ocr.factory.OcrModelFactory;
import cn.smartjavaai.ocr.model.common.detect.OcrCommonDetModel;
import cn.smartjavaai.ocr.model.common.direction.OcrDirectionModel;
import cn.smartjavaai.ocr.model.common.recognize.OcrCommonRecModel;
import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.io.IOUtils;
import org.apache.spark.sql.api.java.UDF1;

import java.io.BufferedInputStream;
import java.io.InputStream;
@Slf4j
public class OcrUDF implements UDF1<byte[], String> {
    public static final String basePath = BaseConstants.basePath;

    public static final DeviceEnum device = BaseConstants.device;

    public static OcrCommonRecModel recModel;

    static {
        log.info("start loading ocr models---------------------------------------------------");
        OcrCommonDetModel detectionModel = getDetectionModel();
        recModel = getRecModel(detectionModel);
        log.info("ocr models loaded---------------------------------------------------");
    }

    public String call(byte[] bytes) throws Exception {
        return extractTextFromImage(bytes);
    }

    public static String  extractTextFromImage(byte[] bytes) {
        try {
            //不带方向矫正，分行返回文本
            OcrRecOptions options = new OcrRecOptions(false, true);
            OcrInfo ocrInfo = recModel.recognize(bytes,options);
            log.info("OCR识别结果：{}", JSONObject.toJSONString(ocrInfo));
            return ocrInfo.getFullText();
        } catch (Exception e) {
            log.error("OCR识别失败",e);
            return null;
        }
    }


    public static OcrCommonRecModel getRecModel(OcrCommonDetModel detModel){
        OcrRecModelConfig recModelConfig = new OcrRecModelConfig();
        //指定文本识别模型
        recModelConfig.setRecModelEnum(CommonRecModelEnum.PP_OCR_V4_MOBILE_REC_MODEL);
        //指定识别模型位置，需要更改为自己的模型路径（下载地址请查看文档）
        recModelConfig.setRecModelPath(basePath+"/ocrmodels/ch_PP-OCRv4_rec_mobile_infer.onnx");
        recModelConfig.setDevice(device);
        recModelConfig.setTextDetModel(detModel);
        return OcrModelFactory.getInstance().getRecModel(recModelConfig);
    }



    /**
     * 获取文本检测模型
     * @return
     */

    public static OcrCommonDetModel getDetectionModel() {
        OcrDetModelConfig config = new OcrDetModelConfig();
        //指定检测模型
        config.setModelEnum(CommonDetModelEnum.PP_OCR_V4_MOBILE_DET_MODEL);
        //指定模型位置，需要更改为自己的模型路径（下载地址请查看文档）
        config.setDetModelPath(basePath+"/ocrmodels/ch_PP-OCRv4_det_mobile_infer.onnx");
        config.setDevice(device);
        return OcrModelFactory.getInstance().getDetModel(config);
    }

    /**
     * 获取方向检测模型
     * @return
     */

    public static   OcrDirectionModel getDirectionModel(){
        DirectionModelConfig directionModelConfig = new DirectionModelConfig();
        //指定行文本方向检测模型
        directionModelConfig.setModelEnum(DirectionModelEnum.PP_LCNET_X0_25);
        //指定行文本方向检测模型路径，需要更改为自己的模型路径（下载地址请查看文档）
        directionModelConfig.setModelPath(basePath+"/ocrmodels/PP-LCNet_x0_25_textline_ori_infer.onnx");
        directionModelConfig.setDevice(device);
        return OcrModelFactory.getInstance().getDirectionModel(directionModelConfig);
    }
}
